#ifndef ROOTEDTREE_H_
#define ROOTEDTREE_H_

#include "common.h"

namespace masugo
{

template<typename VertexKeyType, typename VertexValueType, typename EdgeValueType>
class RootedTree
{
public:
	typedef VertexKeyType VertexKey;
	typedef VertexValueType VertexValue;
	typedef EdgeValueType EdgeValue;

	struct Vertex
	{
		Vertex(const VertexKey parent, const EdgeValue& edge_value, const VertexValue& body_value, const std::size_t level) :
			parent(parent),
			edge_value(edge_value),
			body_value(body_value),
			level(level),
			order(0)
		{
		}

		VertexKey parent;
		EdgeValue edge_value;
		VertexValue body_value;
		std::size_t level;
		int order;
		std::deque< std::pair<int, VertexKey> > children_keys;
	};

	typedef std::tr1::unordered_map<VertexKey, Vertex> Vertices;

	typedef std::tr1::unordered_map<VertexKey, VertexKey> EquivalenceMapping;

	RootedTree(const VertexKey root_key, const VertexValue& root_value) : root_key_(root_key)
	{
		vertices_.insert(std::make_pair(root_key_, Vertex(root_key_, EdgeValue(), root_value, 0)));
		levels_.push_back(std::deque<VertexKey>(1, root_key_));
	}

	VertexKey root_key() const
	{
		return root_key_;
	}

	const Vertices& vertices() const
	{
		return vertices_;
	}

	bool vertex_exists(const VertexKey& vertex_key) const
	{
		return (vertices_.find(vertex_key)!=vertices_.end());
	}

	const Vertex& const_vertex(const VertexKey vertex_key) const
	{
		typename Vertices::const_iterator vertex_it=vertices_.find(vertex_key);
		if(vertex_it==vertices_.end()) { throw std::invalid_argument(utilities::ParametricString("Function {} did not find vertex {}").arg(__FUNCTION__).arg(vertex_key).str()); }
		return vertex_it->second;
	}

	const EquivalenceMapping& const_equivalence_mapping() const
	{
		return equivalence_mapping_;
	}

	const std::string canonical_string() const
	{
		return canonical_string_;
	}

	VertexKey add_child(const VertexKey parent_key, const VertexKey child_key, const EdgeValue& edge_value, const VertexValue& body_value)
	{
		Vertex& parent_vertex=mutable_vertex(parent_key);

		if(vertices_.find(child_key)!=vertices_.end()) { throw std::invalid_argument(utilities::ParametricString("Function {} can't add already existing vertex {}").arg(__FUNCTION__).arg(child_key).str()); }

		Vertex& child_vertex=vertices_.insert(std::make_pair(child_key, Vertex(parent_key, edge_value, body_value, parent_vertex.level+1))).first->second;
		parent_vertex.children_keys.push_back(std::make_pair(child_vertex.order, child_key));

		assert(child_vertex.level<=levels_.size());
		if(child_vertex.level==levels_.size()) { levels_.push_back(std::deque<VertexKey>()); }
		levels_[child_vertex.level].push_back(child_key);

		return child_key;
	}

	void order()
	{
		for(int j=(static_cast<int>(levels_.size())-1);j>=0;j--)
		{
			std::deque<VertexKey>& level=levels_[j];
			assert(!level.empty());
			for(std::size_t i=0;i<level.size();i++)
			{
				sort_vertex_children_by_order(level[i]);
			}
			std::sort(level.begin(), level.end(), lexicographical_less(*this));
			int order=0;
			mutable_vertex(level[0]).order=order;
			for(std::size_t i=1;i<level.size();i++)
			{
				if(!lexicographical_equal(level[i-1], level[i])) { order++; }
				mutable_vertex(level[i]).order=order;
			}
		}
		equivalence_mapping_=construct_equivalence_mapping();
		canonical_string_=construct_breadth_first_canonical_string();
	}

	std::deque<VertexKey> unique_frontier() const
	{
		std::deque<VertexKey> result;
		for(typename EquivalenceMapping::const_iterator it=equivalence_mapping_.begin();it!=equivalence_mapping_.end();it++)
		{
			if(it->first!=utilities::virtual_value<VertexKey>() && it->first==it->second) { result.push_back(it->first); }
		}
		return result;
	}

	std::deque<VertexKey> reduced_unique_frontier() const
	{
		std::deque<VertexKey> result;
		for(typename EquivalenceMapping::const_iterator it=equivalence_mapping_.begin();it!=equivalence_mapping_.end();it++)
		{
			if(it->first!=utilities::virtual_value<VertexKey>() && it->first==it->second && (levels_.size()-const_vertex(it->first).level<3)) { result.push_back(it->first); }
		}
		return result;
	}

private:
	struct lexicographical_less
	{
		typedef RootedTree<VertexKey, VertexValue, EdgeValue> RTree;
		const RTree& rt;
		lexicographical_less(const RTree& rt) : rt(rt) {}

		bool operator () (const VertexKey vertex_a_key, const VertexKey vertex_b_key) const
		{
			const Vertex& vertex_a=rt.const_vertex(vertex_a_key);
			const Vertex& vertex_b=rt.const_vertex(vertex_b_key);
			if(vertex_a.edge_value<vertex_b.edge_value) { return true; }
			else if(vertex_a.edge_value>vertex_b.edge_value) { return false; }
			else if(vertex_a.body_value<vertex_b.body_value) { return true; }
			else if(vertex_a.body_value>vertex_b.body_value) { return false; }
			else
			{
				for(std::size_t i=0;i<std::min(vertex_a.children_keys.size(), vertex_b.children_keys.size());i++)
				{
					if(vertex_a.children_keys[i].first<vertex_b.children_keys[i].first) { return true; }
					else if(vertex_a.children_keys[i].first>vertex_b.children_keys[i].first) { return false; }
				}
				return (vertex_a.children_keys.size()>vertex_b.children_keys.size());
			}
		}
	};

	Vertex& mutable_vertex(const VertexKey vertex_key)
	{
		typename Vertices::iterator vertex_it=vertices_.find(vertex_key);
		if(vertex_it==vertices_.end()) { throw std::invalid_argument(utilities::ParametricString("Function {} did not find vertex {}").arg(__FUNCTION__).arg(vertex_key).str()); }
		return vertex_it->second;
	}

	void sort_vertex_children_by_order(const VertexKey vertex_key)
	{
		Vertex& vertex=mutable_vertex(vertex_key);
		if(!vertex.children_keys.empty())
		{
			for(std::size_t i=0;i<vertex.children_keys.size();i++)
			{
				vertex.children_keys[i].first=const_vertex(vertex.children_keys[i].second).order;
			}
			std::sort(vertex.children_keys.begin(), vertex.children_keys.end());
		}
	}

	bool lexicographical_equal(const VertexKey vertex_a_key, const VertexKey vertex_b_key) const
	{
		const Vertex& vertex_a=const_vertex(vertex_a_key);
		const Vertex& vertex_b=const_vertex(vertex_b_key);
		if(vertex_a.edge_value!=vertex_b.edge_value) { return false; }
		else if(vertex_a.body_value!=vertex_b.body_value) { return false; }
		else if(vertex_a.children_keys.size()!=vertex_b.children_keys.size()) { return false; }
		else
		{
			for(std::size_t i=0;i<std::min(vertex_a.children_keys.size(), vertex_b.children_keys.size());i++)
			{
				if(vertex_a.children_keys[i].first!=vertex_b.children_keys[i].first) { return false; }
			}
			return true;
		}
	}

	EquivalenceMapping construct_equivalence_mapping() const
	{
		EquivalenceMapping equivalence_mapping;
		equivalence_mapping[root_key_]=root_key_;
		for(std::size_t j=1;j<levels_.size();j++)
		{
			const std::deque<VertexKey>& level=levels_[j];
			typedef std::map< std::pair<VertexKey, int>, std::set<VertexKey> > LevelEquivalenceGroups;
			LevelEquivalenceGroups level_equivalence_groups;
			for(std::size_t i=0;i<level.size();i++)
			{
				const Vertex& vertex=const_vertex(level[i]);
				level_equivalence_groups[std::make_pair(equivalence_mapping.find(vertex.parent)->second, vertex.order)].insert(level[i]);
			}
			for(typename LevelEquivalenceGroups::const_iterator it=level_equivalence_groups.begin();it!=level_equivalence_groups.end();it++)
			{
				const std::set<VertexKey>& members=it->second;
				VertexKey smallest_member=(*members.begin());
				for(typename std::set<VertexKey>::const_iterator jt=members.begin();jt!=members.end();jt++)
				{
					equivalence_mapping[*jt]=smallest_member;
				}
			}
		}
		return equivalence_mapping;
	}

	std::string construct_depth_first_canonical_string() const
	{
		const VertexKey backtrack=utilities::null_value<VertexKey>();
		std::ostringstream output;
		std::deque<VertexKey> keys(1, root_key_);
		while(!keys.empty())
		{
			const VertexKey key=keys.back();
			keys.pop_back();
			if(key==backtrack) { output << "$"; }
			else
			{
				const Vertex& vertex=const_vertex(key);
				if(key!=root_key_) { keys.push_back(backtrack); }
				for(typename std::deque< std::pair<int, VertexKey> >::const_reverse_iterator rit=vertex.children_keys.rbegin();rit!=vertex.children_keys.rend();rit++)
				{
					keys.push_back(rit->second);
				}
				output << "(" << vertex.edge_value << "," << vertex.body_value << ")";
			}
		}
		output << "#";
		return output.str();
	}

	std::string construct_breadth_first_canonical_string() const
	{
		const VertexKey backtrack=utilities::null_value<VertexKey>();
		std::ostringstream output;
		std::deque<VertexKey> keys(1, root_key_);
		keys.push_back(backtrack);
		for(std::size_t i=0;i<keys.size();i++)
		{
			const VertexKey key=keys[i];
			if(key==backtrack) { output << "$"; }
			else
			{
				const Vertex& vertex=const_vertex(key);
				for(typename std::deque< std::pair<int, VertexKey> >::const_iterator it=vertex.children_keys.begin();it!=vertex.children_keys.end();it++)
				{
					keys.push_back(it->second);
				}
				keys.push_back(backtrack);
				output << "(" << vertex.edge_value << "," << vertex.body_value << ")";
			}
		}
		output << "#";
		return output.str();
	}

	VertexKey root_key_;
	Vertices vertices_;
	std::deque< std::deque<VertexKey> > levels_;
	EquivalenceMapping equivalence_mapping_;
	std::string canonical_string_;
};

}

#endif /* ROOTEDTREE_H_ */
